package io.github.zerobyteword.metricspringstartup;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.core.metrics.ApplicationStartup;
import org.springframework.core.metrics.StartupStep;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;

public class MetricApplicationStartup implements ApplicationStartup {
    private static final Log logger = LogFactory.getLog(MetricApplicationStartup.class);
    private final AtomicLong idSeq = new AtomicLong(0);

    private final AtomicReference<MetricStartupStep> current = new AtomicReference<>();
    private List<MetricStartupStep> metricStartupSteps = new ArrayList<>();


    @Override
    public StartupStep start(String name) {
        long id = this.idSeq.getAndIncrement();
        long start = System.currentTimeMillis();
        while (true) {
            MetricStartupStep current = this.current.get();
            MetricStartupStep parent = getLatestActive(current);
            MetricStartupStep next = new MetricStartupStep(parent, name, id, start, this::record);
            if (this.current.compareAndSet(current, next)) {
                return next;
            }
        }
    }


    private MetricStartupStep getLatestActive(MetricStartupStep step) {
        while (step != null && step.isEnded()) {
            step = step.getParent();
        }
        return step;
    }


    private void record(MetricStartupStep step) {
        List<MetricStartupStep> startupSteps = this.metricStartupSteps;
        if (startupSteps != null) {
            startupSteps.add(step);
        }
        while (true) {
            MetricStartupStep current = this.current.get();
            MetricStartupStep next = getLatestActive(current);
            if (this.current.compareAndSet(current, next)) {
                return;
            }
        }
    }

    public List<MetricStartupStep> getMetricStartupSteps() {
        return this.metricStartupSteps;
    }


    public void printStartUpTimeCost() {
        List<MetricStartupStep> startupSteps = this.metricStartupSteps;
        if (startupSteps != null) {
            startupSteps.stream()
                    .sorted(Comparator.comparing(MetricStartupStep::costTime).reversed())
                    .forEach(e -> {
                        StringBuilder tagBuilder = new StringBuilder();
                        e.getTags().forEach(v -> tagBuilder.append(v.getKey()).append("=").append(v.getValue()).append(","));
                        logger.info(e.getName() + " " + tagBuilder.toString() + " 耗时:" + e.costTime() + "ms");

                    });
            this.metricStartupSteps = null;
        }
    }
}
